import os
import cv2
import copy
import numpy as np

def computeSIFT(img):
    if len(img.shape)==3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    image = copy.deepcopy(img)
    h, w = image.shape
    image = cv2.resize(image, (h//2, w//2))
    sift = cv2.SIFT_create()
    keypoints_roi, descriptor_roi = sift.detectAndCompute(image, None)
    return keypoints_roi, descriptor_roi

def convert(img_path, txt_path):
    img = cv2.imread(img_path)
    h, w, _ = img.shape
    lists = list()
    with open(txt_path, 'r') as filer:
        for i, line in enumerate(filer.readlines()):
            if i < 5:
                lines = [float(x) for x in line.strip().split(' ') if float(x)>0]
                Cx, Cy, W, H = int(w*lines[0]), int(h*lines[1]), int(w*lines[2]), int(h*lines[3])
                xmin, xmax = int(Cx-W/2), int(Cx+W/2)
                ymin, ymax = int(Cy-H/2), int(Cy+H/2)

                xmin = 0 if xmin<0 else xmin
                xmax = w-1 if xmax>w else xmax
                ymin = 0 if ymin<0 else ymin
                ymax = h-1 if ymax>h else ymax
                lists.append([xmin, ymin, xmax, ymax])
                cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0,0,255), 3)

    return img, lists

def parserKeyPoint(keypoints):
    points = []
    for i, kp in enumerate(keypoints):
        tmp = [kp.pt, kp.size, kp.angle, kp.response, kp.octave, kp.class_id]
        points.append(tmp)
    return points

def generateKeyPoint(points):
    keypoints = []
    for i, point in enumerate(points):
        temp = cv2.KeyPoint(x=point[0][0], y=point[0][1], size=point[1], angle=point[2], \
                response=point[3], octave=point[4], class_id=point[5])
        keypoints.append(temp)
    return keypoints

def generate(img_path, save_path):
    img_lists = [x for x in os.listdir(img_path) if x.endswith('jpg')]

    for i, img in enumerate(img_lists):
        split = img.split('.')
        img_name = split[0]+split[1]+split[2]
        #img_name = img.strip('.jpg')
        attr = dict()
        save_dir = os.path.join(save_path, img_name+'.npy')
        result, lists = convert(os.path.join(img_path, img), os.path.join(img_path, img.replace('.jpg', '.txt')))
        attr['position'] = lists

        image = cv2.imread(os.path.join(img_path, img))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        keypoints, descriptors = computeSIFT(image)
        keypoint = parserKeyPoint(keypoints)

        attr['keypoints'] = keypoint
        attr['descriptors'] = descriptors
        np.save(save_dir, attr)
        cv2.imwrite(save_dir.replace('npy', 'jpg'), result)
        print('Processing index: ', i, ', image: ', img)
    return keypoints

if __name__ == '__main__':
    img_path = '/media/Harddisk/Users/fan/PROJECT/CDPROJECT/TensorRT/Templates'
    save_path = '/media/Harddisk/Users/fan/PROJECT/CDPROJECT/TensorRT/Models/templates'
    keypoint_g = generate(img_path=img_path, save_path=save_path)
    # path = os.path.join(save_path, 'C59FR-L10.npy')
    # data = np.load(path, allow_pickle=True).item()
    # keypoint = data['keypoints']
    # keypoint_p = generateKeyPoint(keypoint)
    print('over')



